package com.cui.project.recommend.util;

import com.huaban.analysis.jieba.JiebaSegmenter;

import java.util.*;
import java.util.concurrent.ConcurrentHashMap;
import java.util.stream.Collectors;

/**
 * TF-IDF + 余弦相似（中文按字级切分）
 *
 * 运行逻辑：
 * 1）fit(docs)：对每个文档进行：
 *    - 分词（英文按词、中文逐字）
 *    - 统计 DF（包含该词的文档数）
 *    - 计算 TF-IDF，并对向量做 L2 归一化
 * 2）topKSimilar(id, k)：
 *    - 取 id 对应向量，与其他文档向量求余弦相似
 *    - 按相似度降序取 Top-K
 *
 * 分析用了那些算法：
 *      TF-IDF 向量化
 *      计算每个接口文档（name/description/...）的 TF-IDF 权重向量。
 *      关键式：tfidf(t,d) = tf(t,d) * log((N+1)/df(t))，并做 L2 归一化。
 *      余弦相似度（Cosine Similarity）
 *      用于计算接口文本之间的内容相似：cos(a,b)= (a·b)/(|a||b|)。
 *      简易中文分词：汉字逐字切分（英文按词），保证 JDK8 纯内存实现。
 */
public class TfidfUtil {

    // 单个文档结构：id + 文本
    public static class Doc {
        public final long id; // 文档ID（此处即接口ID）
        public final String text; // 文本（name + description + tags）
        public Doc(long id, String text){
            this.id = id; // 赋值ID
            this.text = text == null ? "" : text.toLowerCase(); // 文本小写化，null→空
        }
    }

    private final Map<String,Integer> df = new ConcurrentHashMap<String, Integer>(); // 词项 -> 文档频次 DF
    private final Map<Long,Map<String,Double>> tfidf = new ConcurrentHashMap<Long, Map<String, Double>>(); // 文档ID -> (词项 -> TF-IDF)
    private int docCount = 0; // 语料总文档数

    // 训练：构建 DF + TFIDF 向量
    public void fit(List<Doc> docs){
        docCount = docs.size(); // 总文档数
        df.clear(); // 清空 DF
        tfidf.clear(); // 清空向量库

        // —— 先统计 DF（每个文档去重后计一次）——
        for (Doc d: docs){
            Set<String> seen = new HashSet<String>(tok(d.text)); // 文档去重后的词集合
            for (String t: seen){ // 遍历词
                Integer old = df.get(t); // 旧DF
                df.put(t, old == null ? 1 : old + 1); // DF+1
            }
        }
        // —— 计算每个文档的 TF-IDF 向量（并 L2 归一）——
        for (Doc d: docs){
            List<String> toks = tok(d.text); // 分词序列
            // 统计该文档内 TF（词频）
            Map<String,Long> tf = new HashMap<String, Long>(); // 词 -> 频次
            for (String s: toks){
                Long old = tf.get(s); // 旧频次
                tf.put(s, old == null ? 1L : old + 1L); // +1
            }
            Map<String,Double> v = new HashMap<String, Double>(); // 词 -> TF-IDF 值
            for (Map.Entry<String,Long> e: tf.entrySet()){ // 遍历词频
                double tfv = e.getValue() * 1.0 / Math.max(1, toks.size()); // 词频（归一到[0,1]）
                int dfv = Math.max(1, getDf(e.getKey())); // 取该词 DF（>=1 防止除0）
                double idf = Math.log((docCount + 1.0) / dfv); // IDF（加1平滑）
                v.put(e.getKey(), tfv * idf); // TF-IDF 值
            }
            tfidf.put(d.id, norm(v)); // L2 归一并存入“向量库”
        }
    }

    // 查询：给定文档ID，返回 TopK 最相似文档及其相似度
    public Map<Long,Double> topKSimilar(long id, int k){
        Map<String,Double> a = tfidf.get(id); // 取目标向量
        if (a == null || a.isEmpty()) return Collections.<Long,Double>emptyMap(); // 若不存在，返回空
        Map<Long,Double> out = new HashMap<Long, Double>(); // 结果表：文档ID -> 相似度
        for (Map.Entry<Long,Map<String,Double>> e: tfidf.entrySet()){ // 遍历全部文档
            if (e.getKey() == id) continue; // 跳过自身
            out.put(e.getKey(), cos(a, e.getValue())); // 计算余弦相似度
        }
        // 按值倒序取前 k
        List<Map.Entry<Long,Double>> list = new ArrayList<Map.Entry<Long, Double>>(out.entrySet()); // 转列表
        Collections.sort(list, new Comparator<Map.Entry<Long, Double>>() { // 降序
            @Override public int compare(Map.Entry<Long, Double> o1, Map.Entry<Long, Double> o2) {
                return Double.compare(o2.getValue(), o1.getValue());
            }
        });
        LinkedHashMap<Long,Double> res = new LinkedHashMap<Long, Double>(); // 有序 TopK
        int i = 0; // 计数
        for (Map.Entry<Long,Double> e: list){ // 迭代排序后结果
            if (i++ >= k) break; // 超出k停止
            res.put(e.getKey(), e.getValue()); // 写入
        }
        return res; // 返回 TopK 相似结果
    }

    // 辅助：读取 DF 值
    private int getDf(String term){
        Integer v = df.get(term); // DF
        return v == null ? 0 : v.intValue(); // null→0
    }

    // 余弦相似： (a·b)/(|a||b|)
    private static double cos(Map<String,Double> a, Map<String,Double> b){
        if (a == null || b == null || a.isEmpty() || b.isEmpty()) return 0.0; // 任一为空返回0
        double dot = 0.0; // 点积
        for (Map.Entry<String,Double> e: a.entrySet()){ // 遍历 a 的每个维度
            Double y = b.get(e.getKey()); // 取 b 的同维值
            if (y != null) dot += e.getValue() * y; // 累加到点积
        }
        double na = 0.0, nb = 0.0; // 各自范数平方
        for (double x: a.values()) na += x*x; // a 的平方和
        for (double x: b.values()) nb += x*x; // b 的平方和
        na = Math.sqrt(na); nb = Math.sqrt(nb); // 开根号得到范数
        return dot / (na*nb + 1e-12); // 归一化并避免除0
    }

    // L2 归一：将向量各维除以范数
    private static Map<String,Double> norm(Map<String,Double> v){
        double n = 0.0; // 范数平方
        for (double x: v.values()) n += x*x; // 累加平方
        n = Math.sqrt(n); // 开根号
        if (n == 0) return v; // 全0向量直接返回
        Map<String,Double> r = new HashMap<String, Double>(v.size()); // 新map
        for (Map.Entry<String,Double> e: v.entrySet()){ // 遍历每个维度
            r.put(e.getKey(), e.getValue()/n); // 归一
        }
        return r; // 返回归一后向量
    }

    // 简易分词：英文按空格拆分；中文逐字切分（适合轻量场景）
    private static List<String> tok(String s){
        String t = s.replaceAll("[^\\p{IsAlphabetic}\\p{IsDigit}\\u4e00-\\u9fa5]+"," ").trim(); // 过滤非字母/数字/中文，替换为空格
        if (t.isEmpty()) return Collections.<String>emptyList(); // 空串返回空列表
        List<String> out = new ArrayList<String>(); // 结果词表
        String[] arr = t.split("\\s+"); // 空白拆分
        for (String w: arr){ // 遍历每个 token
            //引入分词库（推荐）
            // 添加依赖：com.huaban:jieba-analysis:1.0.2
            if (w.matches(".*[\\u4e00-\\u9fa5].*")) {
                JiebaSegmenter segmenter = new JiebaSegmenter();
                List<String> tokens = segmenter.sentenceProcess(w);
                out.addAll(tokens);
            }
        }
        return out; // 返回词表
    }
}